import json
from http.client import responses

from PyQt5.QtCore import Qt, QObject, pyqtSignal
from PyQt5.QtWidgets import QWidget, QApplication, QVBoxLayout, QHBoxLayout, QPushButton, QScrollArea, QFrame, \
    QComboBox, QLabel, QSlider, QLineEdit, QPlainTextEdit, QFileDialog,QMessageBox
import sys
from PyQt5.QtGui import QIcon, QPixmap
import requests
import threading
import base64

def ai_answer_func(response,signal,ai_answer):
    # 处理流式响应
    try:
        for line in response.iter_lines():
            if line:
                line = line.decode('utf-8')
                if line.startswith('data:'):
                    line = line[5:].strip()
                    if line == '[DONE]':
                        break
                    dic = json.loads(line)
                    if dic.get('type') == 'text':
                        content = dic.get('content')
                        signal.data_updated.emit({'statu':200,'content': content,'output':ai_answer})
            else:
                signal.data_updated.emit({'statu':404})
        signal.data_updated.emit({'statu':100})
    except:
        signal.data_updated.emit({'statu': 500, 'content': '出现了些小问题', 'output': ai_answer})

        signal.data_updated.emit({'statu': 100})



class AI_answer(QObject):
    data_updated = pyqtSignal(dict)

class MainWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.resize(1200, 800)
        self.setWindowTitle("AI Assistant")
        self.setWindowIcon(QIcon("./static/robot.png"))

        # 设置窗口背景色
        self.setStyleSheet("""
            QWidget {
                background-color: #f5f7fa;
                font-family: 'Microsoft YaHei', Arial, sans-serif;
            }
        """)

        main_layout = QVBoxLayout(self)
        main_layout.setContentsMargins(20, 20, 20, 20)
        main_layout.setSpacing(15)

        self.ai_answer_signal=AI_answer()
        self.ai_answer_signal.data_updated.connect(self.update_AI_answer)

        self.net_allow_statu=False
        self.temperature_value=0.5
        self.top_p_value=0.8
        self.history_len=10
        self.history=[]
        self.model='local_model'
        self.content=None
        self.search_ing=False
        self.file_save=False
        self.rag_save=False
        self.answer=''
        self.rag_content=None

        self.initUI(main_layout)

    def initUI(self, main_layout):
        top_layout = QHBoxLayout()
        top_layout.setSpacing(10)
        main_layout.addLayout(top_layout)
        self.create_top_layout(top_layout)

        self.buttom_layout = QHBoxLayout()
        self.buttom_layout.setSpacing(20)
        main_layout.addLayout(self.buttom_layout)
        self.create_chat_layout()

    # 顶部区域
    def create_top_layout(self, top_layout):
        # 添加顶部按钮样式
        button_style = """
            QPushButton {
                background-color: #4a6ee0;
                color: white;
                border: none;
                border-radius: 8px;
                padding: 10px 20px;
                font-size: 14px;
                font-weight: bold;
                min-width: 120px;
            }
            QPushButton:hover {
                background-color: #3a5ec0;
            }
            QPushButton:pressed {
                background-color: #2a4ea0;
            }
        """

        self.chat_btn = QPushButton('聊天')
        self.chat_btn.setStyleSheet(button_style)
        top_layout.addWidget(self.chat_btn)

        self.add_rag_btn = QPushButton('添加知识库')
        self.add_rag_btn.setStyleSheet(button_style)
        top_layout.addWidget(self.add_rag_btn)

        # 添加弹性空间使按钮靠左
        top_layout.addStretch()

        self.chat_btn.clicked.connect(self.create_chat_layout)
        self.add_rag_btn.clicked.connect(self.create_add_rag_layout)

    # 聊天区域
    def create_chat_layout(self):
        self.clear_layout(self.buttom_layout)

        # 左侧设置面板
        left_layout = QVBoxLayout()
        left_layout.setSpacing(15)
        self.buttom_layout.addLayout(left_layout)

        # 右侧聊天区域
        right_layout = QVBoxLayout()
        right_layout.setSpacing(15)
        self.buttom_layout.addLayout(right_layout)

        # 设置左右区域比例
        self.buttom_layout.setStretchFactor(left_layout, 1)
        self.buttom_layout.setStretchFactor(right_layout, 3)

        # 左侧设置面板样式
        left_frame = QFrame()
        left_frame.setFixedWidth(340)
        left_frame.setStyleSheet("""
            QFrame {
                background-color: white;
                border-radius: 12px;
                padding: 20px;
                border: 1px solid #e1e5eb;
            }
        """)
        left_frame_layout = QVBoxLayout(left_frame)
        left_frame_layout.setSpacing(20)
        left_layout.addWidget(left_frame)

        # 模型选择
        model_label = QLabel("选择模型")
        model_label.setStyleSheet("font-weight: bold; color: #2c3e50; font-size: 14px;")
        left_frame_layout.addWidget(model_label)

        self.select_model = QComboBox()
        self.select_model.setStyleSheet("""
            QComboBox {
                background-color: #f8f9fa;
                border: 2px solid #e9ecef;
                border-radius: 8px;
                padding: 8px 12px;
                font-size: 14px;
                min-height: 20px;
            }
            QComboBox::drop-down {
                subcontrol-origin: padding;
                subcontrol-position: top right;
                width: 25px;
                border-left: 1px solid #e9ecef;
            }
            QComboBox::down-arrow {
                image: none;
                border-left: 4px solid transparent;
                border-right: 4px solid transparent;
                border-top: 6px solid #6c757d;
            }
        """)
        self.select_model.addItem('本地模型')
        self.select_model.addItem('网络模型')
        left_frame_layout.addWidget(self.select_model)

        # 温度设置
        temperature_frame = QFrame()
        temperature_layout = QVBoxLayout(temperature_frame)
        temperature_layout.setSpacing(8)
        temperature_header = QHBoxLayout()
        temp_label = QLabel('温度')
        temp_label.setStyleSheet("font-weight: bold; color: #2c3e50; font-size: 14px;")
        temperature_header.addWidget(temp_label)
        temperature_header.addStretch()
        self.temperature_label = QLabel(str(self.temperature_value))
        self.temperature_label.setStyleSheet("color: #4a6ee0; font-weight: bold; font-size: 14px;")
        temperature_header.addWidget(self.temperature_label)
        temperature_layout.addLayout(temperature_header)

        self.temperature_slider = QSlider(Qt.Horizontal)
        self.temperature_slider.setMinimum(0)
        self.temperature_slider.setMaximum(10)
        self.temperature_slider.setValue(int(self.temperature_value*10))
        self.temperature_slider.setStyleSheet("""
            QSlider::groove:horizontal {
                background-color: #e9ecef;
                height: 6px;
                border-radius: 3px;
            }
            QSlider::handle:horizontal {
                background-color: #4a6ee0;
                width: 18px;
                margin: -6px 0;
                border-radius: 9px;
            }
            QSlider::sub-page:horizontal {
                background-color: #4a6ee0;
                border-radius: 3px;
            }
        """)
        temperature_layout.addWidget(self.temperature_slider)
        left_frame_layout.addWidget(temperature_frame)

        # 采样率设置
        top_p_frame = QFrame()
        top_p_layout = QVBoxLayout(top_p_frame)
        top_p_layout.setSpacing(8)

        top_p_header = QHBoxLayout()
        top_p_text = QLabel('采样率')
        top_p_text.setStyleSheet("font-weight: bold; color: #2c3e50; font-size: 14px;")
        top_p_header.addWidget(top_p_text)
        top_p_header.addStretch()
        self.top_p_label = QLabel(str(self.top_p_value))
        self.top_p_label.setStyleSheet("color: #4a6ee0; font-weight: bold; font-size: 14px;")
        top_p_header.addWidget(self.top_p_label)
        top_p_layout.addLayout(top_p_header)

        self.top_p_slider = QSlider(Qt.Horizontal)
        top_p=int(self.top_p_value*100)
        self.top_p_slider.setMinimum(50)
        self.top_p_slider.setMaximum(99)
        self.top_p_slider.setValue(top_p)
        self.top_p_slider.setStyleSheet(self.temperature_slider.styleSheet())
        top_p_layout.addWidget(self.top_p_slider)
        left_frame_layout.addWidget(top_p_frame)

        # 清空历史记录按钮
        clear_chat_history_btn = QPushButton('清空历史记录')
        clear_chat_history_btn.setStyleSheet("""
            QPushButton {
                background-color: #6c757d;
                color: white;
                border: none;
                border-radius: 8px;
                padding: 12px;
                font-size: 14px;
                font-weight: bold;
            }
            QPushButton:hover {
                background-color: #5a6268;
            }
            QPushButton:pressed {
                background-color: #545b62;
            }
        """)
        left_frame_layout.addWidget(clear_chat_history_btn)

        # 添加弹性空间使设置项靠上
        left_frame_layout.addStretch()

        # 右侧聊天区域
        right_frame = QFrame()
        right_frame.setStyleSheet("""
            QFrame {
                background-color: white;
                border-radius: 12px;
                border: 1px solid #e1e5eb;
            }
        """)
        right_frame_layout = QVBoxLayout(right_frame)
        right_frame_layout.setContentsMargins(0, 0, 0, 0)
        right_layout.addWidget(right_frame)

        # 聊天历史区域
        history_frame = self.create_scroll(right_frame_layout)
        self.history_layout = QVBoxLayout(history_frame)
        self.history_layout.setSpacing(10)
        self.history_layout.setContentsMargins(20, 20, 20, 20)

        # 输入区域
        send_frame = QFrame()
        send_frame.setStyleSheet("""
            QFrame {
                background-color: #f8f9fa;
                border-top: 1px solid #e1e5eb;
                padding: 15px;
            }
        """)
        right_frame_layout.addWidget(send_frame)

        input_layout = QVBoxLayout(send_frame)
        input_layout.setSpacing(12)

        # 主输入行
        input_row_layout = QHBoxLayout()
        input_row_layout.setSpacing(10)

        self.input_msg = QLineEdit()
        self.input_msg.setPlaceholderText('输入您的问题...')
        self.input_msg.setStyleSheet("""
            QLineEdit {
                background-color: white;
                border: 2px solid #e9ecef;
                border-radius: 8px;
                padding: 12px 16px;
                font-size: 14px;
            }
            QLineEdit:focus {
                border-color: #4a6ee0;
            }
        """)
        input_row_layout.addWidget(self.input_msg)

        self.send_btn = QPushButton('发送')
        self.send_btn.setStyleSheet("""
            QPushButton {
                background-color: #4a6ee0;
                color: white;
                border: none;
                border-radius: 8px;
                padding: 12px 24px;
                font-size: 14px;
                font-weight: bold;
                min-width: 80px;
            }
            QPushButton:hover {
                background-color: #3a5ec0;
            }
            QPushButton:pressed {
                background-color: #2a4ea0;
            }
        """)
        input_row_layout.addWidget(self.send_btn)
        input_layout.addLayout(input_row_layout)

        # 功能按钮行
        function_layout = QHBoxLayout()
        function_layout.setSpacing(10)

        self.net_allow = QPushButton('联网搜索')
        self.add_file_btn = QPushButton('上传文件')
        self.label_file = QLabel('未上传文件')




        self.net_allow.setStyleSheet('background-color: white;')

        function_layout.addWidget(self.net_allow)
        function_layout.addWidget(self.add_file_btn)
        function_layout.addStretch()
        function_layout.addWidget(self.label_file)

        input_layout.addLayout(function_layout)
        clear_chat_history_btn.clicked.connect(self.clear_chat_history)
        self.select_model.currentTextChanged.connect(self.change_model)
        self.add_file_btn.clicked.connect(self.add_file)
        self.net_allow.clicked.connect(self.net_allow_brn)
        self.send_btn.clicked.connect(self.send_message)
        self.temperature_slider.valueChanged.connect(self.temperature_slider_value_changed)
        self.top_p_slider.valueChanged.connect(self.top_p_slider_value_changed)
    #清空历史记录
    def clear_chat_history(self):
        self.history=[]
        self.clear_layout(self.history_layout)
    #模型类型被更改
    def change_model(self):
        name = self.select_model.currentText()
        if name=='本地模型':
            self.model='local_model'
        else:
            self.model='net_model'
    #温度滑块被修改
    def temperature_slider_value_changed(self):
        self.temperature_value=self.temperature_slider.value()/10
        text=f'{self.temperature_value}'
        self.temperature_label.setText(text)
    #采样率滑块被修改
    def top_p_slider_value_changed(self):
        self.top_p_value = self.top_p_slider.value() / 100
        text = f'{self.top_p_value}'
        self.top_p_label.setText(text)
    #添加文件
    def add_file(self):
        if self.file_save == False:
            try:
                file_filter = "文本文件 (*.txt);;"
                file_paths = QFileDialog.getOpenFileNames(
                    self,
                    '选择多个文件',
                    '../',
                    file_filter
                )
                self.add_file_btn.setText('删除文件')
                with open(file_paths[0][0], 'r') as f:
                   self.content = f.read()
                self.file_save=True
                self.label_file.setText(f'{file_paths[0][0]}')
            except Exception as e:
                print(e)
        else:
            self.content=None
            self.add_file_btn.setText('上传文件')
            self.label_file.setText('未上传文件')
            self.file_save=False
    #联网搜索按钮被点击
    def net_allow_brn(self):
        if self.net_allow_statu:
            self.net_allow_statu=False
            self.net_allow.setStyleSheet('background-color: white;')
        else:
            self.net_allow_statu=True
            self.net_allow.setStyleSheet('background-color: blue;')
    #发送数据给AI
    def send_message(self):
        self.send_btn.setEnabled(False)
        self.send_btn.setText('禁用')
        par = {
            'query': self.input_msg.text(),
            'history': self.history,
            'temperature': self.temperature_value,
            'top_p': self.top_p_value,
            'model': self.model,
            'net_allow_statu': self.net_allow_statu,
            'context': self.content,
        }

        # 创建用户消息框并美化
        user_msg = QPlainTextEdit(f'用户: {self.input_msg.text()}')
        user_msg.setStyleSheet("""
            QPlainTextEdit {
                background-color: #e3f2fd;
                border: 1px solid #90caf9;
                border-radius: 10px;
                padding: 8px;
                font-family: 'Microsoft YaHei';
                font-size: 14px;
                color: #1565c0;
            }
        """)
        user_msg.setReadOnly(True)
        h = len(self.input_msg.text()) // 50 + 1

        user_msg.setFixedHeight(h * 50)
        self.history_layout.addWidget(user_msg)
        self.question=self.input_msg.text()
        # 发送请求
        response = requests.post(url='http://127.0.0.1:8000/api/chat', json=par, stream=True)

        if response.status_code == 200:
            ai_answer = QPlainTextEdit('AI:')
            ai_answer.setStyleSheet("""
                QPlainTextEdit {
                    background-color: #f3e5f5;
                    border: 1px solid #ce93d8;
                    border-radius: 10px;
                    padding: 8px;
                    font-family: 'Microsoft YaHei';
                    font-size: 14px;
                    color: #7b1fa2;
                }
            """)
            ai_answer.setReadOnly(True)
            self.history_layout.addWidget(ai_answer)
            threading.Thread(target=ai_answer_func, args=(response, self.ai_answer_signal, ai_answer)).start()
    # AI回答实时更新
    def update_AI_answer(self,date):
        if date.get('statu')==200:
            content=date.get('content')
            ai_answer=date.get('output')
            old_content = ai_answer.toPlainText()
            text = old_content + content
            h=len(text)//50
            ai_answer.setFixedHeight((h+1)*40)
            ai_answer.setPlainText(text)
            self.answer=self.answer+content
        elif date.get('statu')==100:
            self.send_btn.setText('发送')
            self.send_btn.setEnabled(True)
            self.history.append({'question':self.question,'output':self.answer})
            if len(self.history) > self.history_len:
                self.history.pop(0)
        elif date.get('statu')==500:
            content=date.get('content')
            ai_answer=date.get('output')
            old_content = ai_answer.toPlainText()
            text = old_content + content
            h=len(text)//50
            ai_answer.setFixedHeight((h+1)*40)
            ai_answer.setPlainText(text)
            self.answer=self.answer+content
            self.send_btn.setText('发送')
            self.send_btn.setEnabled(True)
    #添加知识库布局
    def create_add_rag_layout(self):
        self.clear_layout(self.buttom_layout)
        left_layout = QVBoxLayout()
        right_layout=QVBoxLayout()
        self.buttom_layout.addLayout(left_layout)
        self.buttom_layout.addLayout(right_layout)
        rag_name_frame=QFrame()
        left_layout.addWidget(rag_name_frame)
        rag_name_layout=QVBoxLayout(rag_name_frame)
        label=QLabel('知识库名')
        self.rag_name_input=QLineEdit()
        self.rag_name_input.setPlaceholderText('输入知识库名')
        rag_name_layout.addWidget(label)
        rag_name_layout.addWidget(self.rag_name_input)


        rag_des_frame = QFrame()
        left_layout.addWidget(rag_des_frame)
        rag_des_layout = QVBoxLayout(rag_des_frame)
        label = QLabel('知识库描述')
        self.rag_des_input = QPlainTextEdit()
        self.rag_des_input.setPlaceholderText('输入知识库描述')
        rag_des_layout.addWidget(label)
        rag_des_layout.addWidget(self.rag_des_input)

        self.rag_db_frame =self.create_scroll(left_layout)
        layout1 = QVBoxLayout(self.rag_db_frame)
        response=requests.get(url='http://127.0.0.1:8000/api/get_all_db').json()
        data=response.get('data')
        for i in data:
            frame = QFrame()
            frame.setFixedHeight((len(i) // 20 + 1) * 50)
            layout1.addWidget(frame)
            layout2 = QHBoxLayout(frame)
            label = QLabel(f'{i}')
            layout2.addWidget(label)
            show_btn_rag_db=QPushButton('查看')
            show_btn_rag_db.setFixedWidth(80)
            show_btn_rag_db.setProperty('name', i)
            show_btn_rag_db.setStyleSheet("background-color: blue; color: white;")
            show_btn_rag_db.setProperty('name', i)
            layout2.addWidget(show_btn_rag_db)
            delete_btn_rag_db = QPushButton("删除")
            delete_btn_rag_db.setFixedWidth(80)
            delete_btn_rag_db.setStyleSheet("background-color: red; color: white;")
            layout2.addWidget(delete_btn_rag_db)
            delete_btn_rag_db.setProperty('name', i)
            show_btn_rag_db.clicked.connect(self.show_btn_rag_db)
            delete_btn_rag_db.clicked.connect(self.delete_btn_rag_db)
        frame = self.create_scroll(right_layout)
        layout=QVBoxLayout(frame)
        self.rag_text=QPlainTextEdit()
        self.rag_text.setEnabled(False)
        layout.addWidget(self.rag_text)
        add_frame=QFrame()
        right_layout.addWidget(add_frame)
        add_layout = QHBoxLayout()
        right_layout.addLayout(add_layout)
        add_rag_btn=QPushButton('添加知识库')
        add_layout.addWidget(add_rag_btn)
        delete_rag_btn=QPushButton('清空知识库')
        add_layout.addWidget(delete_rag_btn)
        submit_btn=QPushButton('提交知识库')
        add_layout.addWidget(submit_btn)

        add_rag_btn.clicked.connect(self.add_rag)
        delete_rag_btn.clicked.connect(self.delete_rag)
        submit_btn.clicked.connect(self.submit_rag)
    #查看知识库具体信息：
    def show_btn_rag_db(self):
        name=self.sender().property('name')
        par={
            'name':name,
        }
        response=requests.get(url='http://127.0.0.1:8000/api/get_rag_db',params=par).json()
        try:
            data=response.get('data')
            name=data.get('name')
            description=data.get('description')
            text=data.get('text')
            content=f'数据库名：{name}\n\n描述：{description}\n\n{text}'
            self.rag_text.clear()
            self.rag_text.setPlainText(content)
        except:
            return

    #删除知识库文件
    def delete_btn_rag_db(self):
        name=self.sender().property('name')
        par={
            'name':name,
        }
        response=requests.delete(url='http://127.0.0.1:8000/api/drop_rag_db',params=par).json()
        if response.get('code')==200:
            QMessageBox.information(None,'提示','删除成功')
            self.create_add_rag_layout()


    #添加知识库按钮
    def add_rag(self):
        try:
            self.rag_text.clear()
            file_filter = "文本文件 (*.txt);;"
            file_paths = QFileDialog.getOpenFileNames(
                self,
                '选择多个文件',
                '../',
                file_filter
            )
            with open(file_paths[0][0], 'r') as f:
                rag_content = f.read()
            self.rag_save = True
            self.rag_text.setPlainText(rag_content)

        except Exception as e:
            print(e)
    #清空知识库按钮
    def delete_rag(self):
        if self.rag_save == True:
            self.rag_save = False
            self.rag_content=None
            self.rag_text.clear()
    #提交知识库
    def submit_rag(self):
        if self.rag_name_input.text()=='' or self.rag_des_input.toPlainText()=='' or self.rag_text.toPlainText()=='' or self.rag_save==False:
            QMessageBox.critical(None,'提示','请将数据填写完整')
            return
        par={
            'rag_name':self.rag_name_input.text(),
            'rag_description':self.rag_des_input.toPlainText(),
            'rag_text':self.rag_text.toPlainText(),
        }
        response = requests.post(url='http://127.0.0.1:8000/api/add_rag', json=par, stream=True).json()
        if response['code']==200:
            QMessageBox.information(None,'提示','添加成功')
            self.create_add_rag_layout()
        else:
            QMessageBox.information(None,'提示','添加失败')
        self.rag_name_input.clear()
        self.rag_des_input.clear()
        self.rag_text.clear()




    # 创建含滑块的容器
    def create_scroll(self, layout):
        scroll_area = QScrollArea()
        scroll_area.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
        scroll_area.setVerticalScrollBarPolicy(Qt.ScrollBarAsNeeded)
        scroll_area.setWidgetResizable(True)
        scroll_area.setStyleSheet("""
            QScrollArea {
                border: none;
                background-color: transparent;
            }
            QScrollBar:vertical {
                background-color: #f1f3f4;
                width: 10px;
                border-radius: 5px;
            }
            QScrollBar::handle:vertical {
                background-color: #c1c6cc;
                border-radius: 5px;
                min-height: 20px;
            }
            QScrollBar::handle:vertical:hover {
                background-color: #a8acb2;
            }
        """)

        frame = QFrame()
        frame.setStyleSheet("background-color: transparent;")
        scroll_area.setWidget(frame)
        layout.addWidget(scroll_area)
        return frame

    # 清空布局
    def clear_layout(self, layout):
        if layout is None:
            return
        while layout.count():
            item = layout.takeAt(0)
            widget = item.widget()
            if widget:
                widget.deleteLater()
            else:
                self.clear_layout(item.layout())
                item.layout().deleteLater()

    def closeEvent(self, event):
        self.close()


if __name__ == '__main__':
    app = QApplication(sys.argv)
    win = MainWindow()
    win.show()
    sys.exit(app.exec_())